import socket
import subprocess
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import Tuple 


class PortScanner:
    """端口扫描引擎"""
    
    def __init__(self, max_threads: int = 100):
        self.max_threads = max_threads
        self.cancel_scan = False
    
    def set_cancel_flag(self, cancel: bool = True):
        """设置取消扫描标志"""
        self.cancel_scan = cancel
    
    async def scan_host_ports(self, host: str, start_port: int, end_port: int, 
                             progress_callback=None, result_callback=None) -> list:
        """异步扫描主机端口"""
        try:
            # 解析主机名
            target_ip = socket.gethostbyname(host)
        except socket.gaierror:
            raise ValueError("无法解析主机名")
        
        results = []
        
        # 使用线程池进行异步扫描
        with ThreadPoolExecutor(max_workers=self.max_threads) as executor:
            tasks = []
            
            for port in range(start_port, end_port + 1):
                if self.cancel_scan:
                    break
                
                task = asyncio.get_event_loop().run_in_executor(
                    executor, self._check_port, target_ip, port
                )
                tasks.append((port, task))
            
            # 处理结果
            completed_count = 0
            for port, task in tasks:
                if self.cancel_scan:
                    break
                
                try:
                    is_open, service_name, process_info = await task
                    
                    if is_open:
                        result = {
                            'port': port,
                            'status': '开放',
                            'service': service_name,
                            'process': process_info
                        }
                        results.append(result)
                        
                        # 立即回调结果以实时显示
                        if result_callback:
                            result_callback(result)
                    
                    completed_count += 1
                    
                    # 调用进度回调
                    if progress_callback:
                        progress_callback(completed_count, len(tasks), port)
                    
                    # 小延迟以避免UI冻结
                    await asyncio.sleep(0.001)
                
                except Exception as e:
                    print(f"扫描端口 {port} 时出错: {e}")
        
        return results

    async def scan_wsl_ports(self, distro: str, host: str, start_port: int, 
                           end_port: int, progress_callback=None, result_callback=None) -> list:
        """异步扫描WSL端口"""
        try:
            target_ip = socket.gethostbyname(host)
        except socket.gaierror:
            raise ValueError("无法解析主机名")
        
        results = []
        
        with ThreadPoolExecutor(max_workers=self.max_threads) as executor:
            tasks = []
            
            for port in range(start_port, end_port + 1):
                if self.cancel_scan:
                    break
                
                task = asyncio.get_event_loop().run_in_executor(
                    executor, self._check_wsl_port, distro, target_ip, port
                )
                tasks.append((port, task))
            
            completed_count = 0
            for port, task in tasks:
                if self.cancel_scan:
                    break
                
                try:
                    is_open, service_name, process_info = await task
                    
                    if is_open:
                        result = {
                            'port': port,
                            'status': '开放',
                            'service': service_name,
                            'process': process_info
                        }
                        results.append(result)
                        
                        # 立即回调结果以实时显示
                        if result_callback:
                            result_callback(result)
                    
                    completed_count += 1
                    
                    if progress_callback:
                        progress_callback(completed_count, len(tasks), port)
                    
                    await asyncio.sleep(0.001)
                
                except Exception as e:
                    print(f"扫描WSL端口 {port} 时出错: {e}")
        
        return results

    async def scan_docker_ports(self, container: str, start_port: int, 
                              end_port: int, progress_callback=None, result_callback=None) -> list:
        """异步扫描Docker端口"""
        results = []
        
        with ThreadPoolExecutor(max_workers=self.max_threads) as executor:
            tasks = []
            
            for port in range(start_port, end_port + 1):
                if self.cancel_scan:
                    break
                
                task = asyncio.get_event_loop().run_in_executor(
                    executor, self._check_docker_port, container, port
                )
                tasks.append((port, task))
            
            completed_count = 0
            for port, task in tasks:
                if self.cancel_scan:
                    break
                
                try:
                    is_open, service_name, process_info = await task
                    
                    if is_open:
                        result = {
                            'port': port,
                            'status': '开放',
                            'service': service_name,
                            'process': process_info
                        }
                        results.append(result)
                        
                        # 立即回调结果以实时显示
                        if result_callback:
                            result_callback(result)
                    
                    completed_count += 1
                    
                    if progress_callback:
                        progress_callback(completed_count, len(tasks), port)
                    
                    await asyncio.sleep(0.001)
                
                except Exception as e:
                    print(f"扫描Docker端口 {port} 时出错: {e}")
        
        return results
    
    def _check_port(self, ip: str, port: int) -> Tuple[bool, str, str]:
        """检查主机端口"""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1)
            result = sock.connect_ex((ip, port))
            sock.close()
            
            if result == 0:
                from service_mapper import ServiceMapper
                from process_info import ProcessInfo
                
                service_name = ServiceMapper.get_service_name(port)
                process_info = ProcessInfo.get_process_info(port)
                return True, service_name, process_info
            else:
                return False, "", ""
        
        except Exception:
            return False, "", ""
    
    def _check_wsl_port(self, distro: str, ip: str, port: int) -> Tuple[bool, str, str]:
        """检查WSL端口"""
        try:
            cmd = f'wsl -d {distro} netstat -tuln | grep ":{port} "'
            result = subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                encoding="utf-8", errors="ignore"
            )
            
            if result.returncode == 0 and result.stdout.strip():
                from service_mapper import ServiceMapper
                
                service_name = ServiceMapper.get_service_name(port)
                
                # 获取WSL进程信息
                process_cmd = f"wsl -d {distro} lsof -i :{port} 2>/dev/null | tail -n +2"
                process_result = subprocess.run(
                    process_cmd, shell=True, capture_output=True, text=True,
                    encoding="utf-8", errors="ignore"
                )
                
                process_info = "WSL进程"
                if process_result.returncode == 0 and process_result.stdout.strip():
                    lines = process_result.stdout.strip().split("\n")
                    if lines:
                        parts = lines[0].split()
                        if len(parts) >= 2:
                            process_info = f"WSL: {parts[0]} (PID: {parts[1]})"
                
                return True, service_name, process_info
            else:
                return False, "", ""
        
        except Exception:
            return False, "", ""
    
    def _check_docker_port(self, container: str, port: int) -> Tuple[bool, str, str]:
        """检查Docker容器端口"""
        try:
            cmd = f"docker port {container} {port}"
            result = subprocess.run(
                cmd, shell=True, capture_output=True, text=True,
                encoding="utf-8", errors="ignore"
            )
            
            if result.returncode == 0 and result.stdout.strip():
                from service_mapper import ServiceMapper
                
                service_name = ServiceMapper.get_service_name(port)
                
                # 获取容器进程信息
                process_cmd = f'docker exec {container} netstat -tulnp 2>/dev/null | grep ":{port} "'
                process_result = subprocess.run(
                    process_cmd, shell=True, capture_output=True, text=True,
                    encoding="utf-8", errors="ignore"
                )
                
                process_info = f"Docker: {container}"
                if process_result.returncode == 0 and process_result.stdout.strip():
                    lines = process_result.stdout.strip().split("\n")
                    if lines:
                        parts = lines[0].split()
                        if len(parts) >= 7:
                            process_info = f"Docker: {container} - {parts[6]}"
                
                return True, service_name, process_info
            else:
                return False, "", ""
        
        except Exception:
            return False, "", ""